// ++++++++++++++++++++++++++++++++++++++++
// 《零基础Go语言算法实战》源码
// ++++++++++++++++++++++++++++++++++++++++
// Author:廖显东（ShirDon）
// Blog:https://www.shirdon.com/
// Gitee:https://gitee.com/shirdonl/goAlgorithms.git
// Buy link :https://item.jd.com/14101229.html
// ++++++++++++++++++++++++++++++++++++++++

package main

import (
	"fmt"
	"sort"
)

// 定义一个名为 Edge 的结构，它表示图中的加权边。
type Edge struct {
	src  int // 边的源顶点
	dst  int // 边的目标顶点
	cost int // 边缘的成本或重量
}

// Kruskal 算法用于查找连接的加权图的最小生成树。
// 该函数将顶点数 (v) 和边的切片 (es) 作为输入，并返回最小生成树中的最小成本和边。
func Kruskal(v int, es []Edge) (int, []Edge) {
	// 以非递减顺序按权重对边缘进行排序
	sort.Slice(es, func(i, j int) bool {
		return es[i].cost < es[j].cost
	})

	cost := 0              // 将成本初始化为 0
	mst := make([]Edge, 0) // 将最小生成树初始化为空切片
	u := NewSet(v)         // 为顶点创建一个新的 Union-Find 数据结构

	// 遍历排序的边缘
	for _, e := range es {
		// 如果边的源顶点和目的顶点不在同一个集合中，
		// 则将边添加到最小生成树中，合并集合，并更新成本
		if !u.IsSameSet(e.src, e.dst) {
			u.Union(e.src, e.dst)
			mst = append(mst, e)
			cost += e.cost
		}
	}

	// 返回成本和最小生成树
	return cost, mst
}

// 定义一个名为 UnionSet 的结构体，它表示一个 Union-Find 数据结构
type UnionSet struct {
	p []int // 每个顶点的父级
	r []int // 每个顶点的等级
}

// 创建具有给定大小的新 Union-Find 数据结构
func NewSet(size int) *UnionSet {
	uf := new(UnionSet)
	uf.p = make([]int, size) // 将每个顶点的父级初始化为其自身
	uf.r = make([]int, size) // 将每个顶点的等级初始化为 1
	for i := range uf.p {
		uf.p[i] = i
		uf.r[i] = 1
	}
	return uf
}

// 合并包含顶点 a 和 b 的集合
func (uf *UnionSet) Union(a, b int) {
	// 如果 a 的等级大于 b 的等级，则使 b 成为 a 的孩子
	if uf.r[a] > uf.r[b] {
		uf.p[b] = a
	} else {
		// 否则，让 a 成为 b 的孩子
		uf.p[a] = b
		// 如果 a 的等级等于 b 的等级，则增加 a 的等级
		if uf.r[a] == uf.r[b] {
			uf.r[a]++
		}
	}
}

// 找到包含顶点 x 的集合的根
func (uf *UnionSet) Find(x int) int {
	// 如果 x 的父节点不是 x 本身，则递归地找到包含 x 的父节点的集合的根
	if uf.p[x] != x {
		uf.p[x] = uf.Find(uf.p[x])
	}
	return uf.p[x]
}

// 检查顶点 a 和 b 是否在同一集合中
func (uf *UnionSet) IsSameSet(a, b int) bool {
	// 找到包含 a 的集合的根
	a = uf.Find(a)
	b = uf.Find(b)
	return a == b
}

func main() {
	edges := []Edge{{0, 1, 10},
		{0, 2, 6},
		{0, 3, 5},
		{1, 3, 15},
		{2, 3, 4}}
	v := 5

	res1, res2 := Kruskal(v, edges)
	fmt.Println(res1, res2)
}

//$ go run kruskal.go
//19 [{2 3 4} {0 3 5} {0 1 10}]
